# nano_alpha_evolve/program_db.py

import random


class ProgramDB:
    """
    一个简单的内存数据库，用于存储和管理程序及其分数。
    """
    
    def __init__(self, population_size: int = 10, optimization_mode: str = 'max'):
        self.population = []
        self.population_size = population_size
        if optimization_mode not in ['min', 'max']:
            raise ValueError("optimization_mode must be 'min' or 'max'")
        self.optimization_mode = optimization_mode
    
    def _is_better(self, score1, score2):
        """根据优化模式判断 score1 是否优于 score2"""
        if self.optimization_mode == 'max':
            return score1 > score2
        else:  # min
            return score1 < score2
    
    def add_program(self, program: str, score: dict):
        """
        向数据库中添加一个程序，并保持种群规模受限。
        """
        main_score = score.get("main_score")
        if main_score is None:
            return  # 不添加没有分数的程序
        
        if len(self.population) < self.population_size:
            self.population.append((program, score))
        else:
            # 找到最差的程序
            # 如果是最大化模式，最差的就是分数最小的
            # 如果是最小化模式，最差的就是分数最大的
            worst_program_index = 0
            for i in range(1, len(self.population)):
                current_worst_score = self.population[worst_program_index][1].get("main_score")
                candidate_score = self.population[i][1].get("main_score")
                if self._is_better(current_worst_score, candidate_score):
                    worst_program_index = i
            
            worst_score = self.population[worst_program_index][1].get("main_score")
            if self._is_better(main_score, worst_score):
                self.population[worst_program_index] = (program, score)
        
        # 按分数对种群进行排序
        # max 模式下 reverse=True (降序), min 模式下 reverse=False (升序)
        is_reverse = self.optimization_mode == 'max'
        self.population.sort(key=lambda x: x[1].get("main_score", float('-inf')), reverse=is_reverse)
    
    def get_random_program(self):
        """
        从种群中返回一个随机程序。
        """
        return random.choice(self.population)[0]
    
    # --- 添加这个新方法 ---
    def get_random_program_with_score(self):
        """
        从种群中返回一个随机的 (程序, 分数字典) 元组。
        """
        return random.choice(self.population)
    # --------------------
    
    def get_best_programs(self, n: int = 3):
        """
        返回种群中排名前 n 的程序。
        """
        return self.population[:n]
    
    def get_best_program(self):
        """
        返回种群中最好的程序。
        """
        return self.population[0] if self.population else (None, None)